#include "CUDA_vec.h"
#include "CUDA_JacobiRelax.h"
#include "CUDA_full_weighting.h"
#include "CUDA_interpolation.h"
#include "ConjugateGradient.h"
#include <stdio.h>
#include <vector>
#include <cmath>
using namespace std;
struct Level {
	typedef CUDA_vec<float> Vector;
	Vector * xi;
	Vector * bi;

	Vector & x() { return *xi; }
	Vector & b() { return *bi; }

	Level() { xi = NULL; bi = NULL; }
	Level(Vector & x1, Vector & b1):xi(&x1),bi(&b1){}
};


class MultigridSolver {
public:
	typedef CUDA_vec<float> Vector;
	MultigridSolver() {minLevel = 3;}

	void set_min_level(int m) { minLevel = m; }

	int get_iterations(){ return iter; }

	template<typename Matrix>
	void operator()(const Matrix & A, Vector & x, const Vector & b, float tol)
	{
		typedef CUDA_vec<float> Vector;
		typedef CUDA_Full_Weighting RestrictionOp;
		typedef CUDA_Interpolation InterpolationOp;
		typedef CUDA_JacobiRelax RelaxOp;

		int N = (int)sqrt(x.getSize());
		int p = (int)(log(N+1)/log(2));

		Level * levels = new Level[p - minLevel + 1];

		for(int i=0; i <= p - minLevel; i++)
		{
			int N_i = (1 << (p - i) ) - 1;
			print_CUDA_err_msg("Before allocating");
			Vector * xi = new Vector(N_i*N_i); Vector * bi = new Vector(N_i*N_i);
			print_CUDA_err_msg("Before copying");
			if(i==0)
			{
				xi->copyValues(x);
				bi->copyValues(b);
			}
			levels[i].xi = xi;
			levels[i].bi = bi;
		}

		RelaxOp relax_vector(2, 0.8);
		RestrictionOp restrict_vector;
		InterpolationOp interpolate_vector;
		int lowestLevel = p-minLevel;
		iter = 0;
		printf("Initial norm: %f\n", residual_norm(A, levels[0].x(), levels[0].b()));
		//while(residual_norm(A, levels[0].x(), b) > tol)
		for(int j=0; j < 8; j++)
		{
			for(int i=0; i<p - minLevel; i++)
			{
				//printf("Level %d\n", i);
				//printf("Pre-relax norm: %f\n", residual_norm(A, levels[i].x(), levels[i].b()) );
				relax_vector(A, levels[i].x(), levels[i].b());
				//printf("Post-relax norm: %f\n", residual_norm(A, levels[i].x(), levels[i].b()) );

				Vector r(levels[i].x().getSize());
				compute_residual(A, levels[i].x(), levels[i].b(), r);

				restrict_vector(r, levels[i+1].b());
				levels[i+1].x().set_constant(0);
			}
			ConjugateGradient cg;
			cg.solve<Matrix>(A, levels[lowestLevel].x(), levels[lowestLevel].b(), 1e-6);
			//printf("CG\n");
			for(int i=lowestLevel-1; i >= 0; i--)
			{
				Vector r(levels[i].x().getSize());
				interpolate_vector(levels[i+1].x(), r);
				levels[i].x().AXPBY(1, 1, r);
				//printf("Pre-relax norm: %f\n", residual_norm(A, levels[i].x, levels[i].b) );
				relax_vector(A, levels[i].x(), levels[i].b());
				//printf("Post-relax norm: %f\n", residual_norm(A, levels[i].x, levels[i].b) );
			}
			iter++;
			printf("Iter: %d Norm: %f\n", iter, residual_norm(A, levels[0].x(), b));
		}
		x.copyValues(levels[0].x());
		for(int i=0; i<= p - minLevel; i++)
		{
			delete levels[i].xi;
			delete levels[i].bi;
		}
		delete [] levels;
	}
private:
	int minLevel;
	int iter;
};
